/**

This file is part of MaCI/GIMnet.

MaCI/GIMnet is free software: you can redistribute it and/or modify it 
under the terms of the GNU Lesser General Public License as published 
by the Free Software Foundation, either version 3 of the License, or 
(at your option) any later version.

MaCI/GIMnet is distributed in the hope that it will be useful, but WITHOUT 
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 
FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public 
License for more details.

You should have received a copy of the GNU Lesser General Public 
License along with GIMnet. (See COPYING.LESSER) If not, see 
<http://www.gnu.org/licenses/>.

**/
/**
 * \file
 * \brief Simple binary container for communication.
 * \author Antti Maula <antti.maula@tkk.fi>
 */
#include "binbag.h"
#include "owndebug.h"
#include "ownutils.h"
#include <zlib.h>
#include <string.h>
#include <assert.h>

//*****************************************************************************
using namespace gim::binbag;
//*****************************************************************************
#define KBINBAG_HEADER_ID_CONST          0x666
#define KBINBAG_RESIZE_BLOCKSIZE         4096
#define KBINBAG_MAXIMUM_ELEMENTSIZE      16777216
#define KBINBAG_HEADER_ID_CONST_EXT_ZLIB 0x6661   
//*****************************************************************************
unsigned int CBinBag::sBufferResizes = 0;
unsigned int CBinBag::sTotalMovedBytes = 0;
unsigned int CBinBag::sTotalInstanceCopied = 0;
uint16_t CBinBag::KBinBag_Header_ID = KBINBAG_HEADER_ID_CONST;
//*****************************************************************************

CBinBag::CBinBag()
  : iReadOnly(false),
    iIsCompressed(false),
    iBuffer(NULL),
    iBufferSize(0),
    iBufferUsed(0),
    iElementCount(0),
    iElementAddIterator(NULL),
    iElementGetIterator(NULL)
{
  // Verify that Sizeof TBinBlob is correct!
  assert(sizeof(TBinBlob) == SIZEOF_STRUCT_TBINBLOB);
}
//*****************************************************************************

CBinBag::~CBinBag()
{
  /*dPrint(10,"Called.\n\tBufferResizes: %u\n\tTotalMovedBytes: %u\n\tTotalInstanceCopied: %u",
         sBufferResizes, sTotalMovedBytes, sTotalInstanceCopied);
  */
  // If the container is NOT readonly and the iBuffer is set; delete the buffer.
  if (!iReadOnly && iBuffer) {
    delete[] iBuffer;
    iBuffer = NULL;
    iBufferUsed = iBufferSize = 0;
  }
}
//*****************************************************************************

CBinBag::CBinBag(const CBinBag &bb)
  : iReadOnly(false),
    iIsCompressed(false),
    iBuffer(NULL),
    iBufferSize(0),
    iBufferUsed(0),
    iElementCount(0),
    iElementAddIterator(NULL),
    iElementGetIterator(NULL)
{
  // Call (private) assignment operator to carry out operation
  //dPrint(10,"Copy constructor copying BinBag @ %p to BinBag @ %p",
  //       &bb, this);
  *this = bb;
}
//*****************************************************************************

CBinBag &CBinBag::operator=(const CBinBag &aBinBag)
{
  // Call Copy function.
  Copy(aBinBag);
  return *this;
}
//*****************************************************************************

void CBinBag::Copy(const CBinBag &aBinBag)
{
  // Create new, matching size buffer for copy.
  if (iBuffer != NULL) delete[] iBuffer;
  iBuffer = new char [aBinBag.iBufferUsed];
  assert(iBuffer != NULL);

  // Copy the data
  memcpy(iBuffer, aBinBag.iBuffer, aBinBag.iBufferUsed);
    
  // Set values
  iBufferUsed = iBufferSize = aBinBag.iBufferUsed;
  iReadOnly = false; // Set to FALSE, to be freed on EXIT.
  iElementCount = aBinBag.iElementCount;
  iElementAddIterator = (TBinBlob*)(iBuffer+2);
  iElementGetIterator = iElementAddIterator;
  iIsCompressed = aBinBag.iIsCompressed;

  // Update stats:
  // This is copied instance, so increment counter
  ++sTotalInstanceCopied;

  // memcpy was called for all data, increment bytecounter
  sTotalMovedBytes += aBinBag.iBufferUsed;
}
//*****************************************************************************

bool CBinBag::CompressZLIB(const int &aLevel)
{
  if (iIsCompressed) {
    dPrint(ODERROR, "Container is already compressed, will not compress again!");
    return false;
  }
  if (!iBuffer || !iBufferUsed) {
    dPrint(ODERROR, "No data to compress - will not compress!");
    return false;
  }
  
  // Carry out compression
  uint8_t *buf = NULL;
  unsigned int bufsize = 0;
  //dPrint(ODVINFO,"Compression requested, compressing %u bytes with ZLIB", iBufferUsed);

  if (DoCompressZLIB(buf, bufsize, 
                     (uint8_t*)(iBuffer + 2), (iBufferUsed - 2), 
                     aLevel)) {
    //dPrint(ODVINFO,"Compression compressed, %u bytes -> %u bytes", iBufferUsed, bufsize);

    if (iBuffer && !iReadOnly) delete[] iBuffer;
    iBuffer = (char*)buf;
    iBufferUsed = bufsize;
    iBufferSize = bufsize;
    iIsCompressed = true;
    iReadOnly = false;

    // Update Header
    uint16_t *hdrid = (uint16_t *)iBuffer;
    *hdrid = KBINBAG_HEADER_ID_CONST_EXT_ZLIB;
    iElementAddIterator = (TBinBlob*)(iBuffer+2); // Reset elementiterator

    return true;
  } else {
    dPrint(ODERROR, "Compression failed!");
    return false;

  }
}
//*****************************************************************************

bool CBinBag::Decode(const char *aData, const int aDataSize, bool aCopy)
{
  bool result = false;
  const uint16_t *hdrPtr = (uint16_t *)aData;

  if (aDataSize <= 2) {
    dPrintLCRed(ODERROR,"Attempted to Decode an empty buffer as BinBag container (Size %u bytes)", 
                aDataSize);
    return false;
  }
  assert(aData != NULL);
  
  // Clear out current data. (Leave the memory region alone)
  iBufferUsed = 0;
  iElementCount = 0;

  // Check the policy about data copying. (Copy ONLY if ordinary data,
  // compressed data will be copied anyway)
  if (aCopy && *hdrPtr == KBINBAG_HEADER_ID_CONST) {
    // Then, resize buffer
    ResizeBuffer(aDataSize);
    
    // Take copy of buffer
    memcpy(iBuffer, aData, aDataSize);
    iBufferUsed = aDataSize;
    
  } else {
    // Check whether the buffer is NOT NULL, and NOT readonly, then delete
    if (iBuffer != NULL && !iReadOnly) delete[] iBuffer;

    // Just assign the incoming buffer, mark 'iReadOnly' true.
    iBuffer = (char *)aData;
    iBufferUsed = aDataSize;
    iBufferSize = 0;
    iReadOnly = true;
  }

  int dataSize = 0;
  int payload = 0;

  // Now, check that we have enough data for minimal checks ( Type )
  // Check whether the frame is ZLIB compressed
  if (*hdrPtr == KBINBAG_HEADER_ID_CONST_EXT_ZLIB) {
    //dPrint(ODVINFO,"ZLIB Compressed buffer detected, decompressing %u bytes", 
    //       aDataSize);
    
    uint8_t *buf = NULL;
    unsigned int bufsize = 0;
    if (DoDecompressZLIB(buf, bufsize,
                         (uint8_t*)(aData+2), aDataSize-2)) {
      
      //dPrint(ODVINFO,"Decompression completed. Uncompressed datasize now %u bytes", 
      //       bufsize);
      
      if (iBuffer != NULL && !iReadOnly) delete[] iBuffer;
      iBuffer = (char*)buf;
      iBufferUsed = bufsize;
      iBufferSize = bufsize;
      iReadOnly = false;
      
      // Update header pointer, and update contents! (Type has
      // changed due to decompression)
      uint16_t *newhdrPtr = (uint16_t *)iBuffer;
      *newhdrPtr = KBINBAG_HEADER_ID_CONST;
      hdrPtr = newhdrPtr;


    } else {
      dPrint(ODERROR,"ZLIB inflate() failed to decompress buffer!");
      
    }
  }

  // Check that the frame begins with correct Header_ID
  if (*hdrPtr == KBINBAG_HEADER_ID_CONST) {
    // Take pointer to first element entry
    const char *ptr = iBuffer+2;
    const TBinBlob *eptr = (const TBinBlob *)ptr;
    dataSize = 2; // Set to bufferamount
      
    // Increment the 'dataSize' by size of each element in list
    while (dataSize < iBufferUsed) {
      dataSize += sizeof(TBinBlob) + eptr->size; // Increment 'seen buffer'
	
      iElementCount += 1; // Increment element counter
      payload += eptr->size; // Increment payload counter
	
      ptr += sizeof(TBinBlob) + eptr->size; // Calculate new pointer
      eptr = (const TBinBlob *)ptr; // Store as elementpointer
    } 
      
    // Last, check whether the 'calculated' buffer matches the incoming datasize (atleast enough)
    if (dataSize == iBufferUsed) {
      // All OK, frame matches.
      //dPrint(10,"BinBag verified Ok. Contains %d elements, total of %d bytes of payload. (Total size %d bytes)",
      //       iElementCount, payload, dataSize);

      // Last step, Rewind() to allow access.
      result = Rewind();

    } else if (dataSize >= iBufferUsed) {
      dPrint(ODINFO, "Warning; Excessive buffer for BinBag (Got %d bytes, header says %d bytes required)",
             iBufferUsed, dataSize);

      // Last step, Rewind() to allow easy access.
      result = Rewind();
    }
  }

  
  return result;
}
//*****************************************************************************

const char *CBinBag::GetBuffer(void) const
{
  return iBuffer;
}
//*****************************************************************************

int CBinBag::GetBufferSize() const
{
  return iBufferUsed;
}
//*****************************************************************************

void CBinBag::TrimBuffer(void)
{  
  // Do we need to trim?
  if (iBufferSize > iBufferUsed) {
    iBufferSize = iBufferUsed;
    char *newbuffer = new char [iBufferSize];
    
    // Got buffer? Copy & Delete it.
    if (iBuffer != NULL) { 
      // Got data? Copy it.
      if (iBufferUsed) {
        memcpy(newbuffer, iBuffer, iBufferUsed);
        sTotalMovedBytes += iBufferUsed; // Increment total moved bytes counter
      }

      // Delete.
      delete[] iBuffer;
      iBuffer = NULL;
      ++sBufferResizes; // Increment buffer resizes counter
    }

    // Assign the newly allocated buffer.
    iBuffer = newbuffer;
  }
}
//*****************************************************************************

const TBinBlob *CBinBag::AddBinary(const int aType, const char *aData, const int aDataSize)
{
  TBinBlob *blob = NULL;

  // Check
  if (aDataSize > KBINBAG_MAXIMUM_ELEMENTSIZE) {
    dPrintLCRed(ODERROR,"Requested to add element larger than MAXIMUM_ELEMENTSIZE! (max: %dB, req: %dB)",
           KBINBAG_MAXIMUM_ELEMENTSIZE, aDataSize);
    
  } else if (iReadOnly) {
    dPrintLCRed(ODERROR,"Buffer is Read-Only, will not modify!");
  
  } else if (iIsCompressed) {
    dPrintLCRed(ODERROR,"Buffer is compressed, will not modify!");

  } else {
    
    // Store the current pointer to Buffer
    const char *buffer_now = iBuffer;

    // Call resizer to verify buffersize
    ResizeBuffer(iBufferUsed + aDataSize + sizeof(TBinBlob));

    // If elementCount is zero, write the Header first.
    if (iElementCount == 0) {
      unsigned short int *hdrid = (unsigned short int *)iBuffer;
      *hdrid = KBinBag_Header_ID;
      iElementAddIterator = (TBinBlob*)(iBuffer+2); // Reset elementiterator
      iBufferUsed = 2; // Two bytes used (Headerbyte size)
    } else {
      
      // IF the base pointer inside BinBag has changed due to Resize, 
      // Recalculate the iElementAddIterator based on the delta value.
      if (buffer_now != iBuffer) {
        iElementAddIterator = (TBinBlob *)(iBuffer + (((char *)iElementAddIterator) - buffer_now));
      }
    }

    // New element header set at current new position
    TBinBlob *new_element = iElementAddIterator; // The next to read is now the elementInterator
    assert(new_element != NULL);

    // Fill in Element
    new_element->type = aType;
    new_element->size = aDataSize;
    memcpy(new_element->data, aData, aDataSize);
    
    // Last but not least; recalculate pointers and counters
    ++iElementCount;
    iBufferUsed += sizeof(TBinBlob) + aDataSize;
    iElementAddIterator = (TBinBlob *)(((char *)iElementAddIterator) + 
                                       sizeof(TBinBlob) + iElementAddIterator->size);

    // Return pointer to newly created element.
    blob = new_element;
  }

  return blob;
}
//*****************************************************************************

void CBinBag::Clear(const bool aFreeMemory)
{
  // Clear out current data.
  iBufferUsed = 0;
  iElementCount = 0;
  
  // Reset the readOnly state
  if (iReadOnly) {
    // Was readonly - cannot touch the buffer -> Must reset buffer too.
    iBuffer = NULL;
    iBufferSize = 0;
    iReadOnly = false;

  } else if (aFreeMemory && iBuffer != NULL) {
    // Not readonly buffer, but freeing the memory was requested.
    delete[] iBuffer;
    iBuffer = NULL;
    iBufferSize = 0;

  }
}
//*****************************************************************************

bool CBinBag::Rewind()
{
  bool result = false;
  
  if (iElementCount > 0) {
    // Set elementGetIterator at begin.
    iElementGetIterator = (TBinBlob *)(iBuffer+2);
    result = true;

  } else {
    iElementGetIterator = NULL;

  }

  return result;
}
//*****************************************************************************

const TBinBlob *CBinBag::FindBinary(const int aType)
{
  const TBinBlob *blob = NULL;

  // Rewind first, then iterate through elements until first key found
  if (Rewind()) {
    do {
      blob = GetNextBinary();
      if (blob && blob->type == (unsigned char)aType) {
        break;
      }
    } while(blob);
  }

  return blob;
}
//*****************************************************************************

const TBinBlob *CBinBag::GetNextBinary()
{
  const TBinBlob *bin = iElementGetIterator;

  if (iElementGetIterator) {
    
    // The returned element is the current ElementIterator
    
    // Now, do some pointer math to update the element iterator
    char *tptr = (char *)iElementGetIterator; // Set temp to current element
    tptr += sizeof(TBinBlob) + iElementGetIterator->size; // Jump by number of payload bytes and header
    
    // Check whether 
    int req_b = tptr - iBuffer;
    if (req_b >= iBufferUsed) {
      iElementGetIterator = NULL;
    } else {
      iElementGetIterator = (TBinBlob *)tptr;
    }
  }
  return bin; 
}
//*****************************************************************************
//*****************************************************************************
//*****************************************************************************
//*****************************************************************************
//*****************************************************************************
//*****************************************************************************

bool CBinBag::DoCompressZLIB(uint8_t * &aOutput, unsigned int &aOutputSize,
                             const uint8_t *aInput, const unsigned int &aInputSize,
                             const int &aLevel)
{
  bool result = true;
  int ret;
  z_stream strm;
  unsigned char *in = (unsigned char *)aInput;
  unsigned int outputBufferSize = (aInputSize << 1);
  unsigned int outputBufferUsed = 0;

  // Allocate memory by plan, but ADD two bytes for the header (avoid
  // unnecessary data copying). Also, mark the 'out' buffer to begin
  // from buffer+2 bytes.
  unsigned char *outbuffer = new unsigned char [outputBufferSize + 2];
  unsigned char *out = outbuffer + 2;

  // Clear header part (Two bytes)
  memset(outbuffer, 0x00, 2);
  
  // Check that level is valid.
  assert(aLevel >= -1 && aLevel <= 9);

  /* allocate inflate state */
  strm.zalloc = Z_NULL;
  strm.zfree = Z_NULL;
  strm.opaque = Z_NULL;
  ret = deflateInit(&strm, aLevel);
  if (ret != Z_OK) {
    delete[] outbuffer;
    dPrint(ODERROR,"Failed to initialize deflate() for zlib!");
    return false;
  }

  // Give input 
  strm.avail_in = aInputSize;
  strm.next_in = in;

  // Mark output buffer
  strm.avail_out = outputBufferSize;
  strm.next_out = out;
  
  // Run deflate
  ret = deflate(&strm, Z_FINISH);
  assert(ret != Z_STREAM_ERROR);
  
  // Calculate used buffer
  outputBufferUsed = outputBufferSize - strm.avail_out;
  
  // Process result; NOTE only STREAM_END processed. We have
  // reserved TWICE the amount of space for the compressed data,
  // hence it should definitely fit in the buffer at once (otherwise
  // it doesn't have to :)
  if (ret == Z_STREAM_END) {
    //dPrint(ODTEST,"Whole input stream is now compressed. Used output buffer %u/%u bytes", 
    //       outputBufferUsed, outputBufferSize);
    
  } else {
    dPrint(ODERROR,"Failed to Compress (ZLIB deflate) input buffer of %u bytes, zlib returned value %d!", 
           aInputSize, ret);
    result = false;
    
  }
  assert(strm.avail_in == 0); 
  
  // Check result
  if (result) {
    aOutput = outbuffer;
    aOutputSize = outputBufferUsed + 2;

  } else {
    delete[] outbuffer;

  }
  
  (void)deflateEnd(&strm);

  return result;
}
//*****************************************************************************

bool CBinBag::DoDecompressZLIB(uint8_t * &aOutput, unsigned int &aOutputSize,
                               const uint8_t *aInput, const unsigned int &aInputSize)
{
  bool result = true;
  int ret;
  z_stream strm;
  unsigned char *in = (unsigned char *)aInput;
  unsigned int outputBufferSize = MAX(aInputSize << 1, 8192);
  unsigned int outputBufferUsed = 0;
  unsigned char *outbuffer = new unsigned char [outputBufferSize + 2];
  unsigned char *out = outbuffer + 2;
  
  /* allocate inflate state */
  strm.zalloc = Z_NULL;
  strm.zfree = Z_NULL;
  strm.opaque = Z_NULL;
  strm.avail_in = 0;
  strm.next_in = Z_NULL;
  ret = inflateInit(&strm);
  if (ret != Z_OK) {
    delete[] outbuffer;
    return false;
  }
  // Give input 
  strm.avail_in = aInputSize;
  strm.next_in = in;

  /* run inflate() on input until output buffer not full */
  do {
    // Mark output buffer
    strm.avail_out = outputBufferSize - outputBufferUsed;
    strm.next_out = out + outputBufferUsed;

    // Run inflate
    ret = inflate(&strm, Z_NO_FLUSH);
    assert(ret != Z_STREAM_ERROR);  /* state not clobbered */

    outputBufferUsed = outputBufferUsed + ((outputBufferSize - outputBufferUsed) - strm.avail_out);

    if (ret == Z_OK) {
      // Check whether new room is required for decompressed data.
      if (outputBufferUsed == outputBufferSize) {
        // More rooom, so allocate bigger buffer, copy already received data and continue
        const unsigned int newsize = outputBufferSize << 1;
        unsigned char *newbuffer = new unsigned char [newsize+2];
        unsigned char *newout = newbuffer + 2;
        memcpy(newout, out, outputBufferSize);
        delete[] outbuffer;
        outbuffer = newbuffer;
        out = newout;
        outputBufferSize = newsize;
        
      } else {
        // Outputbuffer used is less than available space --> all data decompressed.
      }

    } else if (ret == Z_STREAM_END) {
      //dPrint(ODTEST,"Decompression at Stream end. %u b input left, %u b in output buffer, %u b output available", 
      //strm.avail_in, outputBufferUsed, strm.avail_out);
    
    } else {
      result = false;
      
    }
  } while (strm.avail_out == 0);
  
  // Check result
  if (result) {
    aOutput = outbuffer;
    aOutputSize = outputBufferUsed + 2;

  } else {
    delete[] outbuffer;
    
  }
  
  (void)inflateEnd(&strm);
  
  return result;
}
//*****************************************************************************

void CBinBag::ResizeBuffer(const int aSize)
{
  // Check whether we have enough space reserved
  if (iBufferSize < aSize) {

    // Calculate new buffer size (use BLOCKSIZE increments, FIXME?)
    iBufferSize = ((aSize / KBINBAG_RESIZE_BLOCKSIZE) + 1) * KBINBAG_RESIZE_BLOCKSIZE;
    char *newbuffer = new char [iBufferSize];
    
    // Got buffer? Copy & Delete it.
    if (iBuffer != NULL) { 
      // Got data? Copy it.
      if (iBufferUsed) {
        memcpy(newbuffer, iBuffer, iBufferUsed);
        sTotalMovedBytes += iBufferUsed; // Increment total moved bytes counter
      }

      // Delete.
      delete[] iBuffer;
      iBuffer = NULL;
      ++sBufferResizes; // Increment buffer resizes counter
    }

    // Assign the newly allocated buffer.
    iBuffer = newbuffer;
  }
}
//*****************************************************************************
//*****************************************************************************
//*****************************************************************************
